from torch import Tensor, nn
from CNN import cnnNet
from torch.autograd import Variable
import torch
import torch.utils.data as Data
import numpy as np
from MyDataset import MyDataset

device = torch.device("cuda:0")
cnnModel = net = torch.load('./LeNet.pkl', map_location=torch.device(device))
batch_size = 35
loss_list = []
classify_list = []
label_list = []
# 定义损失函数和
loss_func = nn.BCELoss()
torch.set_grad_enabled(False)
cnnModel.eval()
one_one = 0
one_zero = 0
zero_zero = 0
zero_one = 0
# 获取数据集
# myTestDataset = MyDataset("D:/fuxian2/csv/train", "D:/fuxian2/dataset/train")
myTestDataset = MyDataset("F:/PycharmProjects/ML&DL/csv/test", "F:/PycharmProjects/ML&DL/data/test")
# 读取数据集
test_loader = torch.utils.data.DataLoader(
    dataset=myTestDataset,
    batch_size=batch_size,
    shuffle=True
)

for batch_idx, data in enumerate(test_loader):
    label, txtData = data

    txtData = txtData.to(device, torch.float)
    label = label.to(device, torch.float)

    # squeeze()函数的功能是维度压缩。返回一个tensor，其中输入大小为1的所有维都已删除。
    label = torch.squeeze(label)

    classify_pre = cnnModel(txtData)
    classify_pre = torch.squeeze(classify_pre)

    # classify_pre = Tensor.cpu(classify_pre)
    # print('label:' + str(label.item()))
    print('classify_pre: ' + str(classify_pre))
    print('label: ' + str(label))
    classify_pre = Tensor.cpu(classify_pre)

    Classify_len = len(classify_pre)
    for i in range(int(Classify_len)):
        classify_list.append(classify_pre[i].item())
        label_list.append(label[i].item())
        if classify_list[i] >= 0.5:
            classify_list[i] = 1
        else:
            classify_list[i] = 0
    # print('predict:' + str(classify_list))

        if int(label_list[i]) == 1 & int(classify_list[i]) == 1:
            one_one = one_one + 1
        elif int(label_list[i]) == 1 & int(classify_list[i]) == 0:
            one_zero = one_zero + 1
        elif int(label_list[i]) == 0 & int(classify_list[i]) == 1:
            zero_one = zero_one + 1
        elif int(label_list[i]) == 0 & int(classify_list[i]) == 0:
            zero_zero = zero_zero + 1
        print('batch_idx:' + str(batch_idx))
        print('sum:' + str(one_one + one_zero + zero_one + zero_zero))

    '''
    if classify_pre >= 0.5:
        classify_pre = 1
    else:
        classify_pre = 0
    print('predict:' + str(classify_pre))
    if int(label) == 1 & int(classify_pre) == 1:
        one_one = one_one + 1
    elif int(label) == 1 & int(classify_pre) == 0:
        one_zero = one_zero + 1
    elif int(label) == 0 & int(classify_pre) == 1:
        zero_one += 1
    elif int(label) == 0 & int(classify_pre) == 0:
        zero_zero = zero_zero + 1
    print('batch_idx:' + str(batch_idx))
    print('sum:' + str(one_one + one_zero + zero_one + zero_zero))
    '''

print(one_one)
print(one_zero)
print(zero_one)
print(zero_zero)
oz_sum = one_one + one_zero + zero_one + zero_zero
TP = one_one
FP = one_zero
FN = zero_one
TN = zero_zero
Acc = (TP + TN) / (TP + FP + FN + TN)
Pre = TP / (TP + FP)
Recall = TP / (TP + FN)
F1 = 2 * Pre * Recall / (Pre + Recall)

print("准确率: ", Acc)
print("精确率: ", Pre)
print("召回率: ", Recall)
print("F1分数: ", F1)
'''
print('one_one:'+str(one_one/(one_one+one_zero)))
print('one_zero:'+str(one_zero/(one_one+one_zero)))
print('zero_one:'+str(zero_one/(zero_one+zero_zero)))
print('zero_zero:'+str(zero_zero/(zero_zero+zero_one)))
'''
